import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torch.autograd import Variable
from PIL import Image
from torch.nn import functional as F
import cv2
import numpy as np
from PIL import Image
import sys

from ECSSD.dataloader import get_test_loader,get_train_loader

def prec_rec(y_true, y_pred, beta2):

    eps = sys.float_info.epsilon
    tp = torch.sum(y_true * y_pred)
    all_p_pred = torch.sum(y_pred)
    all_p_true = torch.sum(y_true)

    prec = (tp + eps) / (all_p_pred + eps)
    rec = (tp + eps) / (all_p_true + eps)

    return prec, rec

class solver:
    def __init__(self,model,device, batchsize=1, root="./data/",epoch=40,lr=5e-5,wd=0.005):
        self.model = model
        self.device = device
        self.epoch = epoch
        self.lr = lr
        self.wd = wd
        self.optimizer = optim.Adam(self.model.parameters(),lr=self.lr)
        self.loss = torch.nn.BCEWithLogitsLoss()
        self.train_loader = get_train_loader(root,root+"train_pair.lst",batch_size=batchsize)
        self.test_loader = get_test_loader(root,root+"test_pair.lst")
        self.model = self.model.to(self.device)

    def train(self,path):
        for epoch in range(self.epoch):
            total_loss = 0
            self.model.train()
            for step, data in enumerate(self.train_loader):
                img, gt = data
                img = Variable(img).to(self.device)
                gt = Variable(gt).to(self.device)

                self.optimizer.zero_grad()
                out1, out2 = self.model(img)

                loss = self.loss(out1, gt) + out2
                total_loss += loss.item()

                loss.backward()
                self.optimizer.step()

            if path is not None:
                save_path = path+str(epoch)+".pth"
                torch.save(self.model.state_dict(),save_path)

            print('Epoch:{}    loss:{}'.format(epoch, total_loss/len(self.train_loader)))

    def test(self,path=None):
        total_mae = 0
        total_prec = 0
        total_rec = 0
        with torch.no_grad():
            for step, data in enumerate(self.test_loader):
                img, gt, data_path = data
                img = img.to(self.device)
                gt = gt.to(self.device)
                out, _ = self.model(img)

                mae, prec, rec = self.evaluate(out,gt)

                total_mae += mae
                total_prec += prec
                total_rec += rec

                if path is not None:
                    gen_path = path + data_path[0].split('/')[-1]
                    out = out.data.cpu().numpy().squeeze()
                    out = out*255
                    cv2.imwrite(gen_path, out)

            print("gen pic down")

        total_mae = total_mae/len(self.test_loader)
        total_prec = total_prec/len(self.test_loader)
        total_rec = total_rec/len(self.test_loader)

        f_measure = (((1 + 0.3) * total_prec * total_rec) / (0.3 * total_prec + total_rec)).item()

        print("mae:{}".format(total_mae))
        print("F_measure:{}".format(f_measure))

    def evaluate(self, out, gt):
        threshold = 0.5

        mae =  torch.mean(torch.abs(out-gt), dim=(1, 2, 3)).cpu().numpy()

        gt_arr = gt.squeeze().cpu()
        pred_arr = out.squeeze().cpu()

        max_fmeasure = 0
        best_threshold = 0
        for threshold in range(256):
            threshold = threshold / 255.0
            y_pred = (pred_arr >= threshold).float()
            y_true = (gt_arr >= 0.5).float()

            y_true1 = y_true.view(1, -1)
            y_pred1 = y_pred.view(1, -1)

            prec, rec = prec_rec(y_true1, y_pred1, 0.3)
            fmeasure = ((1 + 0.3) * prec * rec) / (0.3 * prec + rec)
            if fmeasure > max_fmeasure:
                max_fmeasure = fmeasure
                best_threshold = threshold

        # Use the best threshold to calculate precision and recall
        y_pred = (pred_arr >= best_threshold).float()
        y_true = (gt_arr >= 0.5).float()
        y_true1 = y_true.view(1, -1)
        y_pred1 = y_pred.view(1, -1)

        prec, rec = prec_rec(y_true1, y_pred1, 0.3)

        return mae, prec, rec
